import time, os
import torch
import numpy as np
from importlib import import_module
from utils_common import get_time_dif, build_class_map, load_MLB
from wrapper import TokenizerWrapper
from parse import args

class CommonPredicter:
    def __init__(self, config, args, mlb, model, tokenizer, idx_to_class):
        self.config = config
        self.args = args
        self.device = config.device
        self.mlb = mlb
        self.tokenizer = tokenizer
        self.idx_to_class = idx_to_class
        self.model = model
        self.model.eval()
        self.model.to(self.device)

    @classmethod
    def init(cls, args):
        args.infer_mode = True
        import_model = import_module('models.' + args.model)
        config = import_model.Config(args)
        if args.MLB_path: config.MLB_path = args.dataset + "/data/" + args.MLB_path


        tokenizer = TokenizerWrapper.load(config, args)
        if args.multi_labels:
            mlb = load_MLB(config)
        else:
            mlb = None
        
        model = import_model.Model(config).to(config.device)
        model.load_state_dict(torch.load(config.save_path))

        _class_to_idx, idx_to_class = build_class_map(config)

        return cls(config=config, mlb=mlb, model=model, tokenizer=tokenizer, 
                    idx_to_class=idx_to_class, args=args)

    def predict(self, texts):
        tokens_idx_all = self.tokenizer(texts, pad_size=self.args.pad_size, return_tensors="pt", args=args)
        if type(tokens_idx_all) in [list, tuple]:
            # fasttext输入格式: (words_line, seq_len, bigram, trigram)
            tokens_idx_all = [item.to(self.device) for item in tokens_idx_all]
        else:
            # textCNN输入格式: (idxes, seq)
            tokens_idx_all = (tokens_idx_all.to(self.device), 'seq_no_use')
        outputs = self.model(tokens_idx_all)
        
        if args.multi_labels:
            outputs = outputs.data.cpu().numpy()
            predicts_all = np.where(outputs > args.threshold, 1, 0)
            predicts_all_names = self.mlb.inverse_transform(predicts_all)
        else:
            predics = torch.max(outputs.data, 1)[1].cpu().numpy().tolist()
            predicts_all_names = [self.idx_to_class[int(p)] for p in predics]
        
        return predicts_all_names


class BertPredicter:
    def __init__(self, config, args, mlb, model, idx_to_class):
        self.config = config
        self.args = args
        self.device = config.device
        self.mlb = mlb
        self.tokenizer = config.tokenizer
        self.idx_to_class = idx_to_class
        self.model = model
        self.model.eval()
        self.model.to(self.device)

    @classmethod
    def init(cls, args):
        assert "bert" in args.model.lower()
        args.infer_mode = True
        import_model = import_module('models.' + args.model)
        config = import_model.Config(args) # 自动load分词器
        
        if args.MLB_path: config.MLB_path = args.dataset + "/data/" + args.MLB_path

        if args.multi_labels:
            mlb = load_MLB(config)
        else:
            mlb = None
        
        model = import_model.Model(config).to(config.device)
        model.load_state_dict(torch.load(config.save_path))

        _class_to_idx, idx_to_class = build_class_map(config)

        return cls(config=config, mlb=mlb, model=model, idx_to_class=idx_to_class, args=args)

    def predict(self, texts):
        encode_res = self.tokenizer.batch_encode_plus(texts, padding=True, max_length=self.args.pad_size, 
                                pad_to_multiple_of=self.args.pad_size, truncation="longest_first",
                                return_tensors="pt")
        encode_res = {k:v.to(self.device) for k,v in encode_res.items()}
        outputs = self.model(encode_res)
        
        if args.multi_labels:
            outputs = outputs.data.cpu().numpy()
            predicts_all = np.where(outputs > args.threshold, 1, 0)
            predicts_all_names = self.mlb.inverse_transform(predicts_all)
        else:
            predics = torch.max(outputs.data, 1)[1].cpu().numpy().tolist()
            predicts_all_names = [self.idx_to_class[int(p)] for p in predics]
        
        return predicts_all_names


if __name__ == "__main__":

    text = ["从跟班到战友 灵兽伴你畅游传奇世界",
            "2011香港电子展 AEE再度发力新品曝光"]
    # predicter = CommonPredicter.init(args)
    bert_predicter = BertPredicter.init(args)
    print(bert_predicter.predict(text))









